from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser
from langchain_core.agents import AgentFinish, AgentAction
from langchain_core.exceptions import OutputParserException
from langchain_core.messages import BaseMessage
import json
from typing import Any, Union

class FunctionAgentOutputParser(OpenAIFunctionsAgentOutputParser):
    def parse(self, llm_output: Any) -> Union[AgentFinish, AgentAction]:
        message = llm_output.generations[0].message  # 注意这里结构可能略有不同
        if not hasattr(message, "tool_calls") or not message.tool_calls:
            return AgentFinish(
                return_values={"output": message.content or ""},
                log=message.content or ""
            )

        tool_call = message.tool_calls[0]
        name = tool_call["name"]
        args = tool_call["args"]

        try:
            parsed_args = json.loads(args) if isinstance(args, str) else args
        except Exception:
            try:
                fixed = args.replace("'", '"')
                parsed_args = json.loads(fixed)
            except Exception as e:
                raise OutputParserException(f"工具参数 JSON 解析失败: {args}") from e

        return AgentAction(
            tool=name,
            tool_input=parsed_args,
            log=message.content or ""
        )
